# -*- coding: utf-8 -*-
"""
@author:  yinbing.li(yinbing.li@zillnk.com)
@version: 1.0.0
@file:    leetcode542_01矩阵.py
@time:    2023/3/10 18:03
"""
import collections
from typing import List


class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        n, m = len(mat), len(mat[0])
        ans = [[0] * m for _ in range(n)]
        queue = collections.deque()

        # // 找出首轮与0相邻的1，放入队列
        for i in range(n):
            for k in range(m):
                if mat[i][k] == 0:
                    dirs = [[i - 1, k], [i + 1, k], [i, k + 1], [i, k - 1]]
                    for d in dirs:
                        r, c = d
                        if 0 <= r < n and 0 <= c < m and mat[r][c] == 1 and ans[r][c] == 0:
                            queue.appendleft(d)
                            # // 与0相邻的1距离必定为1
                            ans[r][c] = 1

        while len(queue) != 0:
            i, k = queue.pop()
            dirs = [[i - 1, k], [i + 1, k], [i, k + 1], [i, k - 1]]
            for d in dirs:
                r, c = d
                if 0 <= r < n and 0 <= c < m and mat[r][c] == 1 and ans[r][c] == 0:
                    # // 下一轮队列遍历的距离值为上一轮的距离值 + 1
                    ans[r][c] = ans[i][k] + 1
                    queue.appendleft(d)
        return ans


s = Solution()
print(s.updateMatrix([[0, 0, 0], [0, 1, 0], [1, 1, 1]]))